from typing import List


class DistanceLimitedPathsExist:

    def __init__(self, n: int, edges: List[List[int]]):
        # ---------- 数据预处理 ----------
        # 计算结点总数：将所有的结点视为叶结点，将所有的边也视为内部结点
        # 叶结点（结点）的ID从0开始，内部结点（边）的ID从n开始
        m = n + len(edges)

        # 依据边的权重升序排序边
        edges.sort(key=lambda x: x[2])

        def get_top(ii):
            """获取当前结点的最远祖先节点"""
            if ancestor[ii] != ii:
                ancestor[ii] = get_top(ancestor[ii])
            return ancestor[ii]

        # 记录每一条边的距离：self.edges[i]=dis
        # 计算每一个结点的父亲结点：father[i]
        # 计算每一个结点的最优祖先节点：ancestor[i]
        self.edges = {}
        father = [-1] * m
        ancestor = [i for i in range(m)]
        for i, (u, v, dis) in enumerate(edges):
            self.edges[i + n] = dis  # 记录当前边的距离
            father[get_top(u)] = father[get_top(v)] = i + n  # 更新两个结点的最远祖先节点的父亲结点
            ancestor[get_top(u)] = ancestor[get_top(v)] = i + n  # 更新两个结点的最远祖先节点的最远祖先结点

        # ---------- 计算递增上爬法的祖先结点 ----------
        # self.jumps[i][j] = 第i个结点的第2^j个祖先结点
        self.jumps = [[v] if v != -1 else [] for v in father]
        for _ in range(15):
            # 更新父结点信息：每次将父结点更新为父结点的父结点（第二次更新即为父结点的祖父节点，以此类推）
            have_father = False
            for i in range(m):
                if father[i] != -1:
                    father[i] = father[father[i]]
                    have_father = True
                else:
                    father[i] = -1
            if not have_father:
                break

            # 将记录更新的父结点：分别其上的第1个结点、第2个结点、第4个结点……
            for i in range(m):
                if father[i] != -1:
                    self.jumps[i].append(father[i])

    def query(self, p: int, q: int, limit: int) -> bool:
        # ---------- 将p和q中更深的结点爬到和另一个结点相同的深度 ----------
        # 当p和q在相同的深度时，则有：len(self.jumps[p]) == len(self.jumps[q])
        depth_p, depth_q = self._get_depth(p), self._get_depth(q)
        if depth_p < depth_q:
            depth_p, depth_q, p, q = depth_q, depth_p, q, p
        p = self._jump_up(p, depth_p - depth_q)

        while self.jumps[p] and p != q:
            # 从上到下寻找第一个不一样的祖先节点：公共祖先节点一定在不一样的祖先节点的上面
            for i in range(len(self.jumps[p]) - 1, -1, -1):
                if self.jumps[p][i] != self.jumps[q][i]:
                    p, q = self.jumps[p][i], self.jumps[q][i]
                    break
            else:
                # 如果所有祖先节点都一样，那么他们的父结点就是最近的公共祖先节点
                p, q = self.jumps[p][0], self.jumps[q][0]
                break

        return p == q and max(self.edges[p], self.edges[q]) < limit

    def _get_depth(self, node):
        """计算结点深度：O(logN)"""
        step = 0
        while self.jumps[node]:
            step += 1 << (len(self.jumps[node]) - 1)
            node = self.jumps[node][-1]
        return step

    def _jump_up(self, node, step):
        """将结点node向上移动step次：O(logS)"""
        while step:
            jump = step.bit_length() - 1
            step -= 1 << jump
            if jump >= len(self.jumps[node]):
                return -1
            node = self.jumps[node][jump]
        return node


if __name__ == "__main__":
    obj = DistanceLimitedPathsExist(6, [[0, 2, 4], [0, 3, 2], [1, 2, 3], [2, 3, 1], [4, 5, 5]])
    print(obj.query(2, 3, 2))  # True
    print(obj.query(1, 3, 3))  # False
    print(obj.query(2, 0, 3))  # True
    print(obj.query(0, 5, 6))  # False
